import dotenv
from langchain.agents import AgentExecutor, create_xml_agent
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from langchain_community.tools import GoogleSerperRun

dotenv.load_dotenv()

class GoogleSerperSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本的谷歌搜索工具"
        "当你想要回答有关问题的时候，可以调用该工具"
        "该工具的输入是搜索查询语句"
    ),
    api_wrapper=GoogleSerperAPIWrapper()
)

dalle = OpenAIDALLEImageGenerationTool(
    name = "openai-dalle",
    api_wrapper = DallEAPIWrapper(model="dall-e-3")
)

tools = [google_serper,dalle]

prompt = ChatPromptTemplate.from_messages([
    ("human", """You are a helpful assistant. Help the user answer any questions.
 
You have access to the following tools:
 
{tools}
 
In order to use a tool, you can use <tool></tool> and <tool_input></tool_input> tags. You will then get back a response in the form <observation></observation>
For example, if you have a tool called 'search' that could run a google search, in order to search for the weather in SF you would respond:
 
<tool>search</tool><tool_input>weather in SF</tool_input>
<observation>64 degrees</observation>
 
When you are done, respond with a final answer between <final_answer></final_answer>. For example:
 
<final_answer>The weather in SF is 64 degrees</final_answer>
 
Begin!
 
Previous Conversation:
{chat_history}
 
Question: {input}
{agent_scratchpad}"""),
])

llm = ChatOpenAI(model="gpt-4o-mini")
agent = create_xml_agent(
    prompt=prompt,
    tools=tools,
    llm=llm
)

agent_excutor = AgentExecutor(agent=agent,tools=tools,verbode=True)
print(agent_excutor.invoke({"input":"马拉松世界记录是多少","chat_history":""}))